import pickle  # 导入 pickle 模块，用于序列化和反序列化 Python 对象
from collections import Counter  # 导入 Counter 模块，用于统计可哈希对象的频率

import tqdm  # 导入 tqdm 模块，用于显示进度条


class TorchVocab(object):
    """
    定义一个词汇表对象，用于将字段(field)数值化。

    属性:
        freqs: 一个 collections.Counter 对象，保存数据中 token 的频率，用于构建词汇表。
        stoi: 一个 collections.defaultdict 实例，将 token 字符串映射到数值标识符。
        itos: 一个 token 字符串列表，通过数值标识符进行索引。
    """

    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],
                 vectors=None, unk_init=None, vectors_cache=None):
        """
        从 collections.Counter 创建一个 Vocab 对象。

        参数:
            counter: collections.Counter 对象，保存数据中每个值的频率。
            max_size: 词汇表的最大大小，如果为 None 则没有最大值。默认为 None。
            min_freq: 将 token 包含在词汇表中的最小频率。小于 1 的值将被设置为 1。默认为 1。
            specials: 特殊 token (例如，padding 或 eos) 的列表，将添加到词汇表的开头，此外还有一个 <unk> token。
                      默认为 ['<pad>']
            vectors: 可用的预训练向量或自定义预训练向量之一 (参见 Vocab.load_vectors);
                     或上述向量的列表
            unk_init (callback): 默认情况下，将词汇表外的词向量初始化为零向量；
                                 可以是任何接受 Tensor 并返回相同大小的 Tensor 的函数。
                                 默认为 torch.Tensor.zero_
            vectors_cache: 缓存向量的目录。默认为 '.vector_cache'
        """
        self.freqs = counter  # 词频计数器
        counter = counter.copy()  # 复制一份计数器，避免修改原始计数器
        min_freq = max(min_freq, 1)  # 最小频率不能小于 1

        self.itos = list(specials)  # 初始化 itos (index to string) 列表，包含特殊 token
        # 构建词汇表时不计算特殊 token 的频率
        for tok in specials:
            del counter[tok]  # 从计数器中删除特殊 token

        max_size = None if max_size is None else max_size + len(self.itos)  # 计算最大词汇表大小

        # 按频率排序，然后按字母顺序排序
        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])
        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)

        for word, freq in words_and_frequencies:
            if freq < min_freq or len(self.itos) == max_size:
                break  # 如果频率小于最小频率或达到最大大小，则停止添加
            self.itos.append(word)  # 将单词添加到 itos 列表

        # stoi 是 itos 的反向字典
        self.stoi = {tok: i for i, tok in enumerate(self.itos)}  # 创建 stoi (string to index) 字典

        self.vectors = None  # 初始化词向量
        if vectors is not None:
            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)  # 加载词向量
        else:
            assert unk_init is None and vectors_cache is None  # 如果没有提供词向量，则 unk_init 和 vectors_cache 必须为 None

    def __eq__(self, other):
        """判断两个TorchVocab对象是否相等"""
        if self.freqs != other.freqs:
            return False  # 如果词频不相等，返回 False
        if self.stoi != other.stoi:
            return False  # 如果 stoi 字典不相等，返回 False
        if self.itos != other.itos:
            return False  # 如果 itos 列表不相等，返回 False
        if self.vectors != other.vectors:
            return False  # 如果词向量不相等，返回 False
        return True  # 所有属性都相等，返回 True

    def __len__(self):
        """返回词汇表的大小"""
        return len(self.itos)

    def vocab_rerank(self):
        """重新对词汇表进行排序"""
        self.stoi = {word: i for i, word in enumerate(self.itos)}  # 重新创建 stoi 字典

    def extend(self, v, sort=False):
        """扩展词汇表"""
        words = sorted(v.itos) if sort else v.itos  # 获取要扩展的词汇表中的单词
        for w in words:
            if w not in self.stoi:  # 如果单词不在当前词汇表中
                self.itos.append(w)  # 将单词添加到 itos 列表
                self.stoi[w] = len(self.itos) - 1  # 将单词添加到 stoi 字典


class Vocab(TorchVocab):
    """
    定义了用于 BERT 的词汇表, 继承自 TorchVocab
    """

    def __init__(self, counter, max_size=None, min_freq=1):
        self.pad_index = 0  # padding token 的索引
        self.unk_index = 1  # unknown token 的索引
        self.eos_index = 2  # end-of-sentence token 的索引
        self.sos_index = 3  # start-of-sentence token 的索引
        self.mask_index = 4  # mask token 的索引
        super().__init__(counter, specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"],  # 初始化特殊 token
                         max_size=max_size, min_freq=min_freq)  # 调用父类的初始化方法

    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:
        """将句子转换为 token 序列 (未使用)"""
        pass  # 未实现

    def from_seq(self, seq, join=False, with_pad=False):
        """将 token 序列转换为句子 (未使用)"""
        pass  # 未实现

    @staticmethod
    def load_vocab(vocab_path: str) -> 'Vocab':
        """
        从文件中加载词汇表。

        参数:
            vocab_path: str, 词汇表文件的路径

        返回:
            Vocab, 加载的词汇表对象
        """
        with open(vocab_path, "rb") as f:
            return pickle.load(f)  # 使用 pickle 加载词汇表对象

    def save_vocab(self, vocab_path):
        """
        将词汇表保存到文件中。

        参数:
            vocab_path: str, 保存词汇表文件的路径
        """
        with open(vocab_path, "wb") as f:
            pickle.dump(self, f)  # 使用 pickle 保存词汇表对象


# Building Vocab with text files
class WordVocab(Vocab):
    """
    用文本文件构建词汇表
    """

    def __init__(self, texts, max_size=None, min_freq=1):
        """
        初始化 WordVocab 对象。

        参数:
            texts: 迭代器, 包含文本数据的迭代器，可以是列表的列表，也可以是字符串的列表。
            max_size: int, 词汇表的最大大小。
            min_freq: int, 词汇表中单词出现的最小频率。
        """
        print("Building Vocab")  # 打印构建词汇表的消息
        counter = Counter()  # 创建一个计数器
        for line in tqdm.tqdm(texts):  # 遍历文本数据
            if isinstance(line, list):
                words = line  # 如果是列表，则直接使用
            else:
                words = line.replace("\n", "").replace("\t", "").split()  # 如果是字符串，则进行分词

            for word in words:
                counter[word] += 1  # 统计每个单词出现的次数
        super().__init__(counter, max_size=max_size, min_freq=min_freq)  # 调用父类的初始化方法

    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
        """
        将句子转换为 token 序列。

        参数:
            sentence: str 或 list, 要转换的句子。
            seq_len: int, 序列的最大长度。
            with_eos: bool, 是否在序列末尾添加 EOS token。
            with_sos: bool, 是否在序列开头添加 SOS token。
            with_len: bool, 是否返回原始序列的长度。

        返回:
            list, 转换后的 token 序列 (如果 with_len=False) 或 元组 (seq, origin_seq_len) (如果 with_len=True)
        """
        if isinstance(sentence, str):
            sentence = sentence.split()  # 如果是字符串，则进行分词

        seq = [self.stoi.get(word, self.unk_index) for word in sentence]  # 将单词转换为索引，如果单词不在词汇表中，则使用 unk_index

        if with_eos:
            seq += [self.eos_index]  # 如果需要，添加 EOS token
        if with_sos:
            seq = [self.sos_index] + seq  # 如果需要，添加 SOS token

        origin_seq_len = len(seq)  # 记录原始序列的长度

        if seq_len is None:
            pass  # 如果没有指定序列长度，则不进行处理
        elif len(seq) <= seq_len:
            seq += [self.pad_index for _ in range(seq_len - len(seq))]  # 如果序列长度小于最大长度，则进行 padding
        else:
            seq = seq[:seq_len]  # 如果序列长度大于最大长度，则进行截断

        return (seq, origin_seq_len) if with_len else seq  # 根据 with_len 返回不同的结果

    def from_seq(self, seq, join=False, with_pad=False):
        """
        将 token 序列转换为句子。

        参数:
            seq: list, 要转换的 token 序列。
            join: bool, 是否将单词连接成字符串。
            with_pad: bool, 是否包含 padding token。

        返回:
            str 或 list, 转换后的句子 (如果 join=True) 或 单词列表 (如果 join=False)
        """
        words = [self.itos[idx]  # 将索引转换为单词
                 if idx < len(self.itos)  # 如果索引在词汇表范围内
                 else "<%d>" % idx  # 如果索引超出词汇表范围，则使用特殊格式
                 for idx in seq  # 遍历 token 序列
                 if not with_pad or idx != self.pad_index]  # 如果不需要 padding 或索引不是 padding token

        return " ".join(words) if join else words  # 根据 join 返回不同的结果

    @staticmethod
    def load_vocab(vocab_path: str) -> 'WordVocab':
        """
        从文件中加载词汇表。

        参数:
            vocab_path: str, 词汇表文件的路径

        返回:
            WordVocab, 加载的词汇表对象
        """
        with open(vocab_path, "rb") as f:
            return pickle.load(f)  # 使用 pickle 加载词汇表对象


def build():
    """
    构建词汇表并保存到文件。
    """
    import argparse  # 导入 argparse 模块，用于解析命令行参数

    parser = argparse.ArgumentParser()  # 创建一个参数解析器
    parser.add_argument("-c", "--corpus_path", required=True, type=str, help="语料库文件路径")  # 添加 corpus_path 参数
    parser.add_argument("-o", "--output_path", required=True, type=str, help="输出文件路径")  # 添加 output_path 参数
    parser.add_argument("-s", "--vocab_size", type=int, default=None, help="词汇表的最大大小")  # 添加 vocab_size 参数
    parser.add_argument("-e", "--encoding", type=str, default="utf-8", help="文件编码格式")  # 添加 encoding 参数
    parser.add_argument("-m", "--min_freq", type=int, default=1, help="词汇表中单词出现的最小频率")  # 添加 min_freq 参数
    args = parser.parse_args()  # 解析命令行参数

    with open(args.corpus_path, "r", encoding=args.encoding) as f:  # 打开语料库文件
        vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)  # 创建 WordVocab 对象

    print("VOCAB SIZE:", len(vocab))  # 打印词汇表大小
    vocab.save_vocab(args.output_path)  # 保存词汇表到文件


if __name__ == "__main__":
    build()  # 如果作为主程序运行，则调用 build 函数
